package com.rbtree;

/**
 * @Author: sy
 * @Date: 2021/3/23 23:30
 * @Description:
 */
public class RbTree {
    private RbTreeNode mRoot;
    private static final boolean RED = false;
    private static final boolean BLACK = true;

    public RbTree() {
        this.mRoot = null;
    }

    /**
     * 获取某个节点的父节点
     * @param node
     * @return
     */
    private RbTreeNode parentOf(RbTreeNode node) {
        return node != null ? node.parent : null;
    }

    /**
     * 获取某个节点的颜色
     * @param node
     * @return
     */
    private boolean colorOf(RbTreeNode node) {
        return node != null ? node.color : BLACK;
    }

    /**
     * 判断某个节点是否是红色
     * @param node
     * @return
     */
    private boolean isRed(RbTreeNode node) {
        return node != null && node.color == RED;
    }

    /**
     * 判断某个节点是否是黑色
     * @param node
     * @return
     */
    private boolean isBlack(RbTreeNode node) {
        return !isRed(node);
    }

    /**
     * 设置某个节点的颜色为黑色
     * @param node
     */
    private void setBlack(RbTreeNode node) {
        if (node != null) {
            node.color = BLACK;
        }
    }

    /**
     * 设置某个节点的颜色为红色
     * @param node
     */
    private void setRed(RbTreeNode node) {
        if (node != null) {
            node.color = RED;
        }
    }

    /**
     * 设置某个节点的父节点
     * @param node
     * @param parent
     */
    private void setParent(RbTreeNode node, RbTreeNode parent) {
        if (node != null) {
            node.parent = parent;
        }
    }

    /**
     * 设置某个节点的颜色
     * @param node
     * @param color
     */
    private void setColor(RbTreeNode node, boolean color) {
        if (node != null) {
            node.color = color;
        }
    }
    
    private void leftRotate(RbTreeNode pNode) {
        RbTreeNode vNode = pNode.right;
        RbTreeNode rNode = vNode.left;
        pNode.right = rNode;

        if (rNode != null) {
            rNode.parent = pNode;
        }

        vNode.parent = pNode.parent;
        if(pNode.parent == null) {
            this.mRoot = vNode;
        } else {
            if (pNode.parent.left == pNode) {
                pNode.parent.left = vNode;
            } else {
                pNode.parent.right = vNode;
            }
        }
        vNode.left = pNode;
        pNode.parent = vNode;
    }

    private void rightRotate(RbTreeNode pNode) {
        RbTreeNode lNode = pNode.left;
        RbTreeNode fNode = lNode.right;
        pNode.left = fNode;
        if (fNode != null) {
            fNode.parent = pNode;
        }
        lNode.parent = pNode.parent;
        if (pNode.parent == null) {
            this.mRoot = lNode;
        } else {
            if (pNode.parent.left == pNode) {
                pNode.parent.left = lNode;
            } else {
                pNode.parent.right = lNode;
            }
        }
        lNode.right = pNode;
        pNode.parent = lNode;
    }

    public boolean insert(RbTreeNode root, RbTreeNode nodeInserted) {
        if (root == null) {
            mRoot = nodeInserted;
            return true;
        }
        if (root.val > nodeInserted.val) {
            if (root.left == null) {
                root.left = nodeInserted;
                nodeInserted.parent = root;
                nodeInserted.color = RED;
                insertFixUp(nodeInserted);
                return true;
            }else {
                return insert(root.left, nodeInserted);
            }
        } else if (root.val < nodeInserted.val) {
            if (root.right == null) {
                root.right = nodeInserted;
                nodeInserted.parent = root;
                nodeInserted.color = RED;
                insertFixUp(nodeInserted);
                return true;
            } else {
                return insert(root.right, nodeInserted);
            }
        }
        return false;
    }

    /**
     * 红黑树插入修正函数
     * @param currentNode 当前修正节点
     */
    private void insertFixUp(RbTreeNode currentNode) {

        // case 3: 插入节点的父节点是红色
        RbTreeNode parent, gparent;
        while ((parent = parentOf(currentNode)) != null && isRed(parent)) {
            gparent = parentOf(parent);
            // case 3.1 若父节点是祖父节点的左孩子
            if (parent == gparent.left) {
                RbTreeNode uncle = gparent.right;
                // case 3.1.1 当前节点的叔叔节点是红色
                if (isRed(uncle)) {
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(gparent);
                    currentNode = gparent;
                    continue;
                }

                // case 3.1.2 当前节点的父节点是红色，叔叔节点是黑色，且当前节点是其父节点的右孩子
                if (isBlack(uncle) && currentNode == parent.right) {
                    RbTreeNode tmp;
                    leftRotate(parent);
                    tmp = parent;
                    parent = currentNode;
                    currentNode = tmp;
                }

                // case 3.1.3 当前节点的叔叔节点是黑色，且当前节点是左孩子
                setBlack(parent);
                setRed(gparent);
                rightRotate(gparent);
            }
            // case 3.2若父节点是祖父节点的右孩子
            else {
                RbTreeNode uncle = gparent.left;
                // case 3.2.1 当前节点的叔叔节点是红色
                if (isRed(uncle)) {
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(gparent);
                    currentNode = gparent;
                    continue;
                }

                // case 3.2.2 当前节点的父节点是红色，叔叔节点是黑色，且当前节点是其父节点的左孩子
                if (isBlack(uncle) && currentNode == parent.left) {
                    RbTreeNode tmp;
                    leftRotate(parent);
                    tmp = parent;
                    parent = currentNode;
                    currentNode = tmp;
                }

                // case 3.2.3 当前节点的叔叔节点是黑色，且当前节点是右孩子
                setBlack(parent);
                setRed(gparent);
                rightRotate(gparent);
            }
        }
        // case 1: 插入的节点是根节点，直接设置为黑色
        if (currentNode == mRoot) {
            setBlack(currentNode);
            return;
        }
        // case 2: 插入的节点的父节点是黑色，什么也不做，因为不会影响平衡
        if (parentOf(currentNode) != null && isBlack(parentOf(currentNode))) {
            return;
        }
    }

    public RbTreeNode search(RbTreeNode root, int value) {
        if (root == null) {
            return null;
        }
        if (root.val > value) {
            return search(root.left, value);
        } else if (root.val < value) {
            return search(root.right, value);
        } else {
            return root;
        }
    }

    public void remove(int value) {
        RbTreeNode node;
        if ((node = search(this.mRoot, value)) != null) {
            remove(node);
        }
    }

    private void remove(RbTreeNode nodeRemoved) {
        // 被删除节点的左右孩子都不为空
        if (nodeRemoved.left != null && nodeRemoved.right != null) {
            removeNodeWithDoubleChild(nodeRemoved);
        }
        // 被删除节点的左孩子不为空
        else if (nodeRemoved.left != null) {
            removeNodeWithLeftChild(nodeRemoved);
        }
        // 被删除节点的右孩子不为空
        else if (nodeRemoved.right != null) {
            removeNodeWithRightChild(nodeRemoved);
        }
        // 被删除节点左右孩子都为空
        else {
            removeNodeWithNoChild(nodeRemoved);
        }
    }

    private void removeNodeWithDoubleChild(RbTreeNode nodeRemoved) {
        // 找中序后继节点，称为替代节点
        RbTreeNode replace = nodeRemoved.right;
        while (replace.left != null) {
            replace = replace.left;
        }
        // 拷贝替代节点的内容到被删除节点
        nodeRemoved.val = replace.val;
        // 继续看替代节点满足哪种情况
        remove(replace);
    }

    private void removeNodeWithLeftChild(RbTreeNode nodeRemoved) {
        RbTreeNode childOfNodeRemoved, parentOfNodeRemoved;
        boolean colorOfNodeRemoved;
        childOfNodeRemoved = nodeRemoved.left;
        parentOfNodeRemoved = nodeRemoved.parent;
        colorOfNodeRemoved = nodeRemoved.color;
        childOfNodeRemoved.parent = parentOfNodeRemoved;

        if (parentOf(nodeRemoved) != null) {
            if (parentOf(nodeRemoved).left == nodeRemoved) {
                parentOf(nodeRemoved).left = childOfNodeRemoved;
            } else {
                parentOf(nodeRemoved).right = childOfNodeRemoved;
            }
        } else {
            this.mRoot = childOfNodeRemoved;
        }

        if (colorOfNodeRemoved == BLACK) {
            removeFixUp(childOfNodeRemoved, parentOfNodeRemoved);
        }
    }

    private void removeNodeWithRightChild(RbTreeNode nodeRemoved) {
        RbTreeNode childOfNodeRemoved, parentOfNodeRemoved;
        boolean colorOfNodeRemoved;
        childOfNodeRemoved = nodeRemoved.right;
        parentOfNodeRemoved = nodeRemoved.parent;
        colorOfNodeRemoved = nodeRemoved.color;
        childOfNodeRemoved.parent = parentOfNodeRemoved;

        if (parentOf(nodeRemoved) != null) {
            if (parentOf(nodeRemoved).left == nodeRemoved) {
                parentOf(nodeRemoved).left = childOfNodeRemoved;
            } else {
                parentOf(nodeRemoved).right = childOfNodeRemoved;
            }
        } else {
           this.mRoot = childOfNodeRemoved;
        }

        if (colorOfNodeRemoved == BLACK) {
            removeFixUp(childOfNodeRemoved, parentOfNodeRemoved);
        }
    }

    private void removeNodeWithNoChild(RbTreeNode nodeRemoved) {
        RbTreeNode childOfNodeRemoved;
        RbTreeNode parentOfNodeRemoved;
        boolean colorOfNodeRemoved;

        childOfNodeRemoved = null;
        parentOfNodeRemoved = nodeRemoved.parent;
        colorOfNodeRemoved = nodeRemoved.color;

        // 被删除节点不是根节点
        if (parentOf(nodeRemoved) != null) {
            if (parentOf(nodeRemoved).left == nodeRemoved) {
                parentOf(nodeRemoved).left = childOfNodeRemoved;
            } else {
                parentOf(nodeRemoved).right = childOfNodeRemoved;
            }
        } else {
            this.mRoot = childOfNodeRemoved;
        }

        // 删除修正
        if (colorOfNodeRemoved == BLACK) {
            removeFixUp(childOfNodeRemoved, parentOfNodeRemoved);
        }
    }

    /**
     * 删除修正
     * @param childOfNodeRemoved
     * @param parentOfNodeRemoved
     * 将当前节点加上额外的一层黑色，为什么？因为删掉的节点是黑色，相当于少了一层黑色，树就不平衡了，所以先补上一层黑色
     * 然后将这多余的黑色往根节点移动，为什么往根节点移动？因为当前节点是位于底部的节点，需要自上而下调整整个树来维持平衡，而且
     * 根节点限定了只能是黑色，往根节点移动更加容易调整
     * 在向上移动过程中，有可能会出现下面几种情况：
     * 1. 当前节点是“红+黑”，直接变成黑色。原因：黑色已经平衡了，红色节点不影响红黑树的平衡
     * 2. 当前节点是“黑+黑”，且当前节点是根节点，直接变成黑色
     * 3. 非前面两种情况
     */
    private void removeFixUp(RbTreeNode childOfNodeRemoved, RbTreeNode parentOfNodeRemoved) {
        // 第三种情况，当前节点是黑+黑节点
        RbTreeNode other;
        while ((childOfNodeRemoved == null || isBlack(childOfNodeRemoved)) && childOfNodeRemoved != this.mRoot) {
            // 如果当前节点是左节点，这关系到叔叔节点的位置
            if (parentOfNodeRemoved.left == childOfNodeRemoved) {
                other = parentOfNodeRemoved.right;
                // 3.1 当前节点是“黑+黑”节点，且当前节点的兄弟节点是红色
                if (isRed(other)) {
                    setBlack(other);
                    setRed(parentOfNodeRemoved);
                    leftRotate(parentOfNodeRemoved);
                    other = parentOfNodeRemoved.right;
                }
                // 3.2 当前节点是“黑+黑”节点，且当前节点的兄弟节点是黑色，且兄弟的两个孩子也是黑色
                if ((other.left == null || isBlack(other.left)) && (other.right == null || isBlack(other.right))) {
                    setRed(other);
                    childOfNodeRemoved = parentOfNodeRemoved;
                    parentOfNodeRemoved = parentOf(childOfNodeRemoved);
                } else {
                    // 3.3 当前节点是“黑+黑”节点，且当前节点的兄弟节点是黑色，且兄弟节点的左孩子是红色，右孩子是黑色
                    if (other.right == null || isBlack(other.right)) {
                        setBlack(other.left);
                        setRed(other);
                        rightRotate(other);
                        other = parentOfNodeRemoved.right;
                        // 转换成状态3.4
                    }
                    // 3.4 当前节点是“黑+黑”节点，且当前节点的兄弟节点是黑色，且兄弟的左孩子颜色任意，右孩子是红色
                    setColor(other, colorOf(parentOfNodeRemoved));
                    setBlack(parentOfNodeRemoved);
                    setBlack(other.right);
                    leftRotate(parentOfNodeRemoved);
                    childOfNodeRemoved = this.mRoot;
                    break;
                }
            } else {
                other = parentOfNodeRemoved.left;
                // 3.1 当前节点是“黑+黑”节点，且当前节点的兄弟节点是红色
                if (isRed(other)) {
                    setBlack(other);
                    setRed(parentOfNodeRemoved);
                    rightRotate(parentOfNodeRemoved);
                    other = parentOfNodeRemoved.left;
                }
                // 3.2 当前节点是“黑+黑”节点，且当前节点的兄弟节点是黑色，且兄弟的两个孩子也是黑色
                if ((other.left == null || isBlack(other.left)) && (other.right == null || isBlack(other.right))) {
                    setRed(other);
                    childOfNodeRemoved = parentOfNodeRemoved;
                    parentOfNodeRemoved = parentOf(childOfNodeRemoved);
                } else {
                    // 3.3 当前节点是“黑+黑”节点，且当前节点的兄弟节点是黑色，且兄弟节点的右孩子是红色，左孩子是黑色
                    if (other.left == null || isBlack(other.left)) {
                        setBlack(other.right);
                        setRed(other);
                        leftRotate(other);
                        other = parentOfNodeRemoved.left;
                        // 转换成状态3.4
                    }
                    // 3.4 当前节点是“黑+黑”节点，且当前节点的兄弟节点是黑色，且兄弟的右孩子颜色任意，左孩子是红色
                    setColor(other, colorOf(parentOfNodeRemoved));
                    setBlack(parentOfNodeRemoved);
                    setBlack(other.left);
                    rightRotate(parentOfNodeRemoved);
                    childOfNodeRemoved = this.mRoot;
                    break;
                }
            }
        }
        // 第一、二种情况
        setBlack(childOfNodeRemoved);
    }
}
